#include <assert.h>
#include "RPC.hpp"
#include "UnitTest.hpp"
#include "Util.hpp"
#include "StreamBuffer.hpp"
#include "Trace.hpp"
#include <set>
#include <stdio.h>

#define STREAM_ESC 16
#define STREAM_LITERALESC 0
#define STREAM_CONTINUE 1
#define STREAM_CLOSE 2

MissingParam::MissingParam(const char *name)
	:std::runtime_error(name)
{
}

NoSuchFunc::NoSuchFunc(const char *name)
	:std::runtime_error(name)
{
}

ProtocolError::ProtocolError(const char *description)
	:std::runtime_error(description)
{
}

/////////////////////////////////////////////////////////////////////////////

RPCParams::RPCParams()
{
}

RPCParams::RPCParams(std::string str)
{
	vector<string> tokens;
	separateArgs(tokens, str.c_str());
	
	for(size_t ii=0; ii<tokens.size(); ii++)
	{
		string token = tokens[ii];
		size_t jj;
		for(jj=0; jj<token.size(); jj++) {
			if(token[jj]=='=')
				break;
		}
		if(jj==token.size())
			continue;
		
		string name = token.substr(0, jj);
		string value = token.substr(jj+1, token.length()-jj-1);
		params[name] = unquoteString(value.c_str());
	}
}

std::string RPCParams::toString() const
{
	std::string ret = "";
	bool first = true;
	
	for(ParamsType::const_iterator ii=params.begin(); ii!=params.end(); ii++)
	{
		if(first) first = false;
		else ret += " ";
		
		ret += ii->first+"="+quoteProtectString(ii->second);
	}
	
	return ret;
}

bool RPCParams::operator==(const RPCParams &rhs) const
{
	if(rhs.params.size() != params.size())
		return false;
	
	for(ParamsType::const_iterator ii=params.begin(); ii!=params.end(); ii++)
	{
		ParamsType::const_iterator rhsIt = rhs.params.find(ii->first);
		
		if(rhsIt == rhs.params.end())
			return false;
		if(ii->second != rhsIt->second)
			return false;
	}
	
	return true;
}

TEST_SET(RPCParams)
{
	try {
		// Construct some simple parameter sets, stringify them, test properties,
		// and verify that their stringified results are as expected
		RPCParams a;
		a.set("foo", 123);
		a.set("bar", "string with spaces");
		testAssert(a.toString() == "bar=\"string with spaces\" foo=123");
		testAssert(a.hasArg("foo"));
		testAssert(a.hasArg("bar"));
		testAssert(!a.hasArg("non"));
		testAssert(a==a);
		
		RPCParams b;
		b.set("foo", 123);
		testAssert(!(a==b));
		b.set("bar", "string with spaces");
		testAssert(a==b);
		b.set("bar", 321);
		testAssert(!(a==b));
		
		RPCParams c("x=1 y=\"a b c\"");
		testAssert(c.getInt("x") == 1);
		testAssert(c.getString("y") == string("a b c"));
		
		RPCParams a_copy(a.toString());
		testAssert(a==a_copy);
		
		RPCParams d;
		d.set("x", "\\");
		d.set("y", "\\\"");
		d.set("z", "\\\"\\");
		testAssert(d.getString("x") == string("\\"));
		testAssert(d.getString("y") == string("\\\""));
		testAssert(d.getString("z") == string("\\\"\\"));
		
		RPCParams d_copy(d.toString());
		testAssert(d_copy.getString("x") == string("\\"));
		testAssert(d_copy.getString("y") == string("\\\""));
		testAssert(d_copy.getString("z") == string("\\\"\\"));
		
	} catch(std::runtime_error e) {
		testAssert(false);
	}
	
	bool exceptionThrown = false;
	RPCParams empty;
	try {
		empty.getInt("x");
	} catch(std::runtime_error e) {
		exceptionThrown = true;
	}
	testAssert(exceptionThrown);
}

void RPCParams::set(const char *name, bool val)
{
	if(val)
		set(name, "1");
	else
		set(name, "0");
}

void RPCParams::set(const char *name, int val)
{
	char buf[16];
	sprintf(buf, "%i", val);
	set(name, buf);
}

void RPCParams::set(const char *name, const char *val)
{
	params[name] = val;
}

void RPCParams::set(const char *name, std::string val)
{
	params[name] = val;
}

void RPCParams::set(const char *name, int64_t val)
{
	char buf[32];
	sprintf(buf, "%lli", (long long int)val);
	set(name, buf);
}

bool RPCParams::hasArg(const char *name) const
{
	return params.find(name) != params.end();
}

bool RPCParams::getBool(const char *name) const throw(MissingParam)
{
	const char *buf = getString(name);
	int ret;
	sscanf(buf, "%i", &ret);
	return ret ? true : false;
}

int RPCParams::getInt(const char *name) const throw(MissingParam)
{
	const char *buf = getString(name);
	int ret;
	sscanf(buf, "%i", &ret);
	return ret;
}

const char *RPCParams::getString(const char *name) const throw(MissingParam)
{
	ParamsType::const_iterator it = params.find(name);
	if(it != params.end())
		return it->second.c_str();
	else
		throw MissingParam(name);
}

int64_t RPCParams::getInt64(const char *name) const throw(MissingParam)
{
	const char *buf = getString(name);
	long long int ret;
	sscanf(buf, "%lli", &ret);
	return ret;
}

/////////////////////////////////////////////////////////////////////////////

static std::map<std::string, RPCFunc> *rpcHandlers;

void rpcAddHandler(const char *name, RPCFunc func)
{
	assert(func != NULL);
	assert(name != NULL);
	if(!rpcHandlers)
	    rpcHandlers = new std::map<std::string, RPCFunc>();
	(*rpcHandlers)[name] = func;
}

void rpcCallLocal(const char *name, const RPCParams *params, RemoteHost *host)
{
	assert(rpcHandlers != NULL);
	std::map<std::string, RPCFunc>::iterator it = rpcHandlers->find(name);
	if(it == rpcHandlers->end()) throw NoSuchFunc(name);
	
	(it->second)(params, host);
}

/////////////////////////////////////////////////////////////////////////////

RemoteHost::RemoteHost()
{
	this->incomingStream = NULL;
	this->maxStreamId = 1;
	this->sendingStreamId = 0;
	
#ifdef THREADED
	pthread_mutex_init(&sendMutex, NULL);
	pthread_mutex_init(&recvMutex, NULL);
	pthread_mutex_init(&streamsMutex, NULL);
#endif
}

RemoteHost::~RemoteHost()
{
#ifdef THREADED
	pthread_mutex_destroy(&sendMutex);
	pthread_mutex_destroy(&recvMutex);
	pthread_mutex_destroy(&streamsMutex);
#endif
}

void RemoteHost::receivedData(int len, const char *data)
{
	RPC_LOCK(recvMutex);
	
	if(incomingData.getLen() > 0) {
		incomingData.append(len, data);
		int bytesHandled = handleData(incomingData.getLen(), incomingData.getData());
		incomingData.popFront(bytesHandled);
	} else {
		int bytesHandled = handleData(len, data);
		if(bytesHandled < len)
			incomingData.append(len-bytesHandled, data+bytesHandled);
	}
	
	RPC_UNLOCK(recvMutex);
}

/////////////////////////////////////////////////////////////////////////////

void RemoteHost::addStreamReceiver(int streamId, StreamReceiver *receiver)
{
	RPC_LOCK(streamsMutex);
		streamReceivers[streamId] = receiver;
	RPC_UNLOCK(streamsMutex);
}

StreamReceiver *RemoteHost::getStreamReceiver(int streamId)
{
	StreamReceiver *ret;
	
	RPC_LOCK(streamsMutex);
		std::map<int,StreamReceiver*>::iterator ii = streamReceivers.find(streamId);
		if(ii != streamReceivers.end())
			ret = ii->second;
		else
			ret = NULL;
	RPC_UNLOCK(streamsMutex);
	
	return ret;
}

void RemoteHost::closeStream(int streamId)
{
	StreamReceiver *receiver = getStreamReceiver(streamId);
	
	if(receiver) {
		receiver->close();
	}
	
	RPC_LOCK(streamsMutex);
		streamReceivers.erase(streamId);
	RPC_UNLOCK(streamsMutex);
}

/////////////////////////////////////////////////////////////////////////////

int RemoteHost::handleData(int len, const char *data)
{
	const char *pos = data;
	const char *end = data+len;
	
	while(pos<end)
	{
		int bytesHandled = handleChunk(pos, end);
		if(bytesHandled > 0) {
			pos += bytesHandled;
		} else {
			break;
		}
	}
	
	return pos-data;
}

int RemoteHost::handleChunk(const char *data, const char *end)
{
	if(incomingStream)
	{
		// Handle either up to the end marker, if there is one. If not, handle
		// all the data here, possibly minus a byte which could be a truncated
		// part of an end marker.
		
		// Walk through looking for an end marker, and also count escaped bytes
		int escapes = 0;
		const char *pos;
		bool endMarked = false;
		for(pos=data; pos<end; pos++)
		{
			if(*pos == STREAM_ESC)
			{
				if(pos+1>=end) {
					break;
				} else if(*(pos+1) == STREAM_LITERALESC) {
					escapes++;
					pos++;
				} else {
					endMarked = true;
					break;
				}
			}
		}
		
		int blockLen = pos-data-escapes;
		
		if(blockLen > 0)
		{
			if(escapes > 0)
			{
				char *unescaped = (char*)malloc(blockLen);
				const char *unescapedReadPos = data;
				for(int ii=0; ii<blockLen; ii++)
				{
					unescaped[ii] = *unescapedReadPos;
					if(*unescapedReadPos == STREAM_ESC) {
						unescapedReadPos+=2;
					} else {
						unescapedReadPos++;
					}
				}
				
				if(incomingStream)
					incomingStream->recv(blockLen, unescaped);
				free(unescaped);
			}
			else
			{
				incomingStream->recv(blockLen, data);
			}
		}
		
		if(endMarked) {
			pos++;
			int marker = *(pos++);
			if(marker==STREAM_CLOSE) {
				incomingStream->close();
			}
			incomingStream = NULL;
		}
		
		return pos-data;
	}
	else
	{
		const char *pos;
		for(pos=data; pos<end; pos++)
		{
			if(*pos == '\n') {
				handleLine(data, pos);
				return pos - data + 1;
			}
		}
		return 0;
	}
}

void RemoteHost::handleLine(const char *line, const char *end)
{
	if(*line == '\x1b')
	{
		// HACK: Skip protocol handshake line
		return;
	}
	else if(isdigit(*line))
	{
		const char *pos;
		for(pos = line; pos < end; pos++)
		{
			if(!isdigit(*pos))
				throw ProtocolError("Line starting with a number has other options");
		}
		
		int streamid = atoi(line);
		incomingStream = getStreamReceiver(streamid);
		if(!incomingStream)
			throw ProtocolError("Data received for closed or invalid stream\n");
	}
	else
	{
		const char *pos;
		for(pos = line; pos < end; pos++)
		{
			if(isspace(*pos))
				break;
		}
		
		string cmd = std::string(line, pos-line);
		RPCParams args(string(pos, end-pos));
		
		rpcCallLocal(cmd.c_str(), &args, this);
	}
}

/////////////////////////////////////////////////////////////////////////////

int RemoteHost::newStreamId()
{
	int ret;
	RPC_LOCK(streamsMutex);
		ret = maxStreamId++;
	RPC_UNLOCK(streamsMutex);
	return ret;
}

/////////////////////////////////////////////////////////////////////////////

void RemoteHost::enterStream(int streamId)
{
	if(sendingStreamId)
	{
		if(streamId == sendingStreamId)
			return;
		else
			exitStream();
	}
	
	char buf[16];
	sprintf(buf, "%i\n", streamId);
	sendData(strlen(buf), buf);
	sendingStreamId = streamId;
}

void RemoteHost::exitStream()
{
	if(!sendingStreamId)
		return;
	
	char buf[2];
	buf[0] = STREAM_ESC;
	buf[1] = STREAM_CONTINUE;
	sendData(2, buf);
	sendingStreamId = 0;
}

void RemoteHost::call(const char *name, const RPCParams *params)
{
	RPC_LOCK(sendMutex);
	
	if(sendingStreamId)
		exitStream();
	
	string sendString = string(name) + " " + params->toString() + "\n";
	sendData(sendString.length(), sendString.c_str());
	
	RPC_UNLOCK(sendMutex);
}

void RemoteHost::sendStream(int streamId, int len, const char *data)
{
	if(!len)
		return;
	
	RPC_LOCK(sendMutex);
	
	enterStream(streamId);
	
	int escapes = 0;
	for(int ii=0; ii<len; ii++) {
		if(data[ii] == STREAM_ESC)
			escapes++;
	}
	
	if(escapes)
	{
		int escapedLen = len+escapes;
		char *escaped = (char*)malloc(escapedLen);
		int writepos = 0;
		for(int ii=0; ii<len; ii++)
		{
			escaped[writepos++] = data[ii];
			if(data[ii] == STREAM_ESC)
				escaped[writepos++] = STREAM_LITERALESC;
		}
		
		sendData(escapedLen, escaped);
		free(escaped);
	}
	else
	{
		sendData(len, data);
	}
	
	RPC_UNLOCK(sendMutex);
}

void RemoteHost::sendCloseStream(int streamId)
{
	RPC_LOCK(sendMutex);
	
	enterStream(streamId);
	char buf[2];
	buf[0] = STREAM_ESC;
	buf[1] = STREAM_CLOSE;
	sendData(2, buf);
	
	sendingStreamId = 0;
	
	RPC_UNLOCK(sendMutex);
}

/////////////////////////////////////////////////////////////////////////////

StreamSender::StreamSender(RemoteHost *destination, RPCParams *params)
{
	this->destination = destination;
	this->streamId = destination->newStreamId();
	this->closed = false;
	
	params->set("stream", streamId);
}

StreamSender::~StreamSender()
{
	if(!closed)
		close();
}

void StreamSender::send(int len, const char *data)
{
	if(closed)
		return;
	
	destination->sendStream(streamId, len, data);
}

void StreamSender::close()
{
	if(!closed) {
		destination->sendCloseStream(streamId);
		this->closed = true;
	}
}

bool StreamSender::isClosed() const
{
	return closed;
}

void StreamSender::disconnected()
{
	this->closed = true;
}

/////////////////////////////////////////////////////////////////////////////

StreamReceiver::StreamReceiver(RemoteHost *source, const RPCParams *params)
{
	this->source = source;
	this->streamId = params->getInt("stream");
	this->closed = false;
	source->addStreamReceiver(streamId, this);
}

StreamReceiver::~StreamReceiver()
{
	close();
}

void StreamReceiver::close()
{
	if(closed)
		return;
	
	this->closed = true;
	onClose();
	source->closeStream(streamId);
}

bool StreamReceiver::isClosed()
{
	return closed;
}

void StreamReceiver::onClose()
{
}

void StreamReceiver::disconnected()
{
	onClose();
}

/////////////////////////////////////////////////////////////////////////////

#ifdef UNIT_TEST

class StreamTestHost
	:public RemoteHost
{
public:
	StreamTestHost(RemoteHost *destination)
	{
		this->destination = destination;
	}
	
	void sendData(int len, const char *data)
	{
		assert(destination);
		
		int pos = 0;
		while(pos < len)
		{
			int chunkSize = randInt(1, 16);
			if(chunkSize+pos > len)
				chunkSize = len-pos;
			
			destination->receivedData(chunkSize, data+pos);
			
			pos += chunkSize;
		}
	}
	
	RemoteHost *destination;
};
class StreamTestReceiver
	:public StreamReceiver
{
public:
	StreamTestReceiver(RemoteHost *source, const RPCParams *params)
		:StreamReceiver(source, params)
	{
		this->received = "";
		this->closeNotified = false;
	}
	void recv(int len, const char *data) {
		received += string(data, len);
	}
	void onClose() {
		closeNotified = true;
	}
	void clear() {
		received = "";
	}

	string received;
	bool closeNotified;
};

struct TestStream
{
	StreamSender *sender;
	StreamTestReceiver *receiver;
};

RPCParams lastReceivedParams;
StreamTestReceiver *lastReceivedStream;

RPC_HANDLER(unittestcmd)
{
	if(params->hasArg("stream")) {
		lastReceivedStream = new StreamTestReceiver(host, params);
	}
	
	lastReceivedParams = *params;
}

string fuzzString()
{
	int len = randInt(0, 15);
	string ret = "";
	for(int ii=0; ii<len; ii++)
	{
		switch(randInt(0,3)) {
			case 0: ret += '\\'; break;
			case 1: ret += '\0'; break;
			case 2: ret += '\"'; break;
			case 3: ret += (char)randInt(0,255); break;
		}
	}
	return ret;
}

RPCParams fuzzParams()
{
	RPCParams ret;
	int numParams = randInt(0, 10);
	for(int ii=0; ii<numParams; ii++)
	{
		char paramName[2] = {'a',0};
		paramName[0] = 'a'+ii;
		ret.set(paramName, fuzzString());
	}
	return ret;
}

TEST_SET(RPC)
{
	std::set<TestStream*> openStreams;
	
	StreamTestHost *receiver = new StreamTestHost(NULL);
	StreamTestHost *sender = new StreamTestHost(receiver);
	
	for(int ii=0; ii<10000; ii++)
	{
		// Randomly either
		//  - Run a command
		//  - Send data on an open stream (if there is one; otherwise open one)
		//  - Open a new stream
		//  - Close one (50%) or all (50%) streams
		switch(randInt(0, 3))
		{
			case 0: {
				RPCParams params = fuzzParams();
				sender->call("unittestcmd", &params);
				testAssert(params == lastReceivedParams);
				break;
			}
				
			case 1: {
				if(openStreams.size())
				{
					// Send data on an open stream
					int sendIndex = randInt(0, openStreams.size()-1);
					std::set<TestStream*>::iterator jj = openStreams.begin();
					for(int kk=0; kk<sendIndex; kk++) jj++;
					
					string data = fuzzString();
					(*jj)->receiver->clear();
					(*jj)->sender->send(data.length(), data.c_str());
					testAssert((*jj)->receiver->received == data);
					break;
				}
			}
			//FALLTHROUGH
			
			case 2: {
				TestStream *stream = new TestStream();
				openStreams.insert(stream);
				RPCParams params;
				stream->sender = new StreamSender(sender, &params);
				sender->call("unittestcmd", &params);
				stream->receiver = lastReceivedStream;;
				testAssert(stream->receiver->received == "");
				break;
			}
				
			case 3: {
				if(!openStreams.size())
					break;
				
				if(randInt(0,1)) {
					// Close one stream at random
					int closeIndex = randInt(0, openStreams.size()-1);
					std::set<TestStream*>::iterator jj = openStreams.begin();
					for(int kk=0; kk<closeIndex; kk++) jj++;
					
					delete (*jj)->sender;
					testAssert((*jj)->receiver->closeNotified);
					delete (*jj)->receiver;
					delete *jj;
					
					openStreams.erase(jj);
				} else {
					// Close all streams
					for(std::set<TestStream*>::iterator jj = openStreams.begin(); jj != openStreams.end(); jj++)
					{
						delete (*jj)->sender;
						testAssert((*jj)->receiver->closeNotified);
						delete (*jj)->receiver;
						delete *jj;
					}
					
					openStreams.clear();
				}
				break;
			}
		}
	}
	
	delete sender;
	delete receiver;
}

#endif

